import torch

class AdversarialGenerator:
    def forward(self, x_real, x_label, ξ=0.11):
        # 这里实现生成扰动的方法
        # 例如：随机生成一个与 x_real 形状相同的扰动
        return torch.randn_like(x_real) * ξ

    def generate_adv_samples(
        self,
        x_real,
        x_label,
        ξ=0.1,
    ):
        # 生成初始扰动
        perturbation: torch.Tensor = self.forward(x_real, x_label, ξ=ξ)
        
        # 对扰动进行裁剪，确保扰动在 [-ξ, ξ] 范围内
        clamped_perturbation = torch.clamp(perturbation, -ξ, ξ)
        
        # 确保所有值为正数
        positive_perturbation = torch.abs(clamped_perturbation)
        
        # 生成对抗样本
        x_adv = x_real + positive_perturbation
        
        # 创建一个掩码，只对原始数据 x 中为 0 或 1 的位置进行四舍五入
        binary_mask = (x_real == 0) | (x_real == 1)  # 仅保留 0 和 1 的特征

        # 对 x_adv 中 binary_mask 选中的值进行四舍五入
        x_adv = torch.where(
            binary_mask,  # 只对 0 或 1 的特征应用四舍五入
            torch.round(x_adv),  # 四舍五入
            x_adv  # 其他特征值保持原样
        )

        return x_adv
# 示例数据（多条样本）
x_real = torch.tensor([
    [0.0, 1.0, 0.5, 0.7],  # 样本 1
    [1.0, 0.0, 0.3, 0.2],  # 样本 2
    [0.0, 0.0, 0.8, 0.9],  # 样本 3
], dtype=torch.float32)  # 形状为 (3, 4)

x_label = torch.tensor([
    [0],  # 样本 1 的标签
    [1],  # 样本 2 的标签
    [0],  # 样本 3 的标签
], dtype=torch.float32)  # 形状为 (3, 1)

# 创建对抗样本生成器
adv_generator = AdversarialGenerator()

# 生成对抗样本
x_adv = adv_generator.generate_adv_samples(x_real, x_label, ξ=0.1)

print("原始样本:")
print(x_real)
print("\n对抗样本:")
print(x_adv)